from faker import Faker
import random
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from data_mock.utils import FileUtil, MysqlUtils_AI

fake = Faker('zh_CN')

# ========== 住院相关基础数据 ==========
DEPARTMENTS = [
    ("D01", "肿瘤内科", "W01", "肿瘤内科一病区"),
    ("D02", "肿瘤外科", "W02", "肿瘤外科一病区"),
    ("D03", "放疗科", "W03", "放射治疗病区"),
    ("D04", "血液科", "W04", "血液病区"),
    ("D05", "中西医结合科", "W05", "中西医结合病区")
]

BED_CODES = [f"B{str(i).zfill(3)}" for i in range(1, 51)]
BED_NAMES = [f"{i}床" for i in range(1, 51)]

MEDICAL_TEAMS = [
    ("MT01", "肿瘤内科一组"),
    ("MT02", "肿瘤内科二组"),
    ("MT03", "肿瘤外科一组"),
    ("MT04", "放疗科一组"),
    ("MT05", "中西医结合组")
]

ADMISSION_TYPES = [
    ("01", "门诊入院"),
    ("02", "急诊入院"),
    ("03", "转科入院"),
    ("04", "转院入院")
]

DISCHARGE_REASONS = [
    "治疗完成出院",
    "病情好转出院",
    "自动出院",
    "转院治疗",
    "病情恶化死亡"
]

FEE_TYPES = [
    "医保",
    "自费",
    "公费",
    "商业保险",
    "新农合"
]

HOSPITAL_INFO = [
    ("H1001", "北京协和医院", "P01", "东院"),
    ("H1002", "上海瑞金医院", "P02", "总院"),
    ("H1003", "广州中山医院", "P03", "肿瘤分院")
]


# ========== 时间范围配置 ==========
def generate_DAY_RANGE(start_day, end_day):
    start = datetime.strptime(start_day, "%Y-%m-%d")
    end = datetime.strptime(end_day, "%Y-%m-%d")
    days = []
    current = start
    while current <= end:
        days.append(current.strftime("%Y-%m-%d"))
        current += relativedelta(days=1)
    return days


DAY_RANGE = generate_DAY_RANGE("2025-08-01", "2025-08-03")


def generate_records_per_day(day: str):
    date_obj = datetime.strptime(day, '%Y-%m-%d')
    return 10 if date_obj.day % 2 == 0 else 6


# ========== 核心生成函数 ==========
def generate_inpatient_records():
    sql_statements = []
    record_count = 0

    patient_id_counter = 2001
    visit_sn_counter = 2001
    inpatient_no_counter = 500001

    for day in DAY_RANGE:
        RECORDS_PER_day = generate_records_per_day(day)
        for i in range(1, RECORDS_PER_day + 1):
            record_count += 1
            patient_id = f"PT{patient_id_counter}"
            patient_id_counter += 1

            visit_sn = f"VIS{visit_sn_counter}"
            visit_sn_counter += 1

            # 随机选择医院信息
            hospital_code, hospital_name, branch_code, branch_name = random.choice(HOSPITAL_INFO)

            # 基础患者信息
            medical_record_no = f"MR{random.randint(100000, 999999)}"
            inpatient_no = f"IP{inpatient_no_counter}"
            inpatient_no_counter += 1
            name = fake.name()

            # 科室和床位信息
            dept_code, dept_name, ward_code, ward_name = random.choice(DEPARTMENTS)
            admission_bed_code = random.choice(BED_CODES)
            admission_bed_name = BED_NAMES[BED_CODES.index(admission_bed_code)]

            # 随机决定是否转科/转床
            if random.random() > 0.7:  # 30%概率转科或转床
                discharge_dept_code, discharge_dept_name, discharge_ward_code, discharge_ward_name = random.choice(
                    DEPARTMENTS)
                discharge_bed_code = random.choice(BED_CODES)
                discharge_bed_name = BED_NAMES[BED_CODES.index(discharge_bed_code)]
            else:
                discharge_dept_code, discharge_dept_name = dept_code, dept_name
                discharge_ward_code, discharge_ward_name = ward_code, ward_name
                discharge_bed_code, discharge_bed_name = admission_bed_code, admission_bed_name

            # 医疗团队信息
            medical_team_code, medical_team_name = random.choice(MEDICAL_TEAMS)

            # 时间信息
            admission_datetime = fake.date_time_between(
                start_date=datetime.strptime(day, "%Y-%m-%d"),
                end_date=(datetime.strptime(day, "%Y-%m-%d") + timedelta(days=1))
            ).strftime('%Y-%m-%d %H:%M:%S')

            # 住院天数3-30天
            hospitalization_days = random.randint(3, 30)
            discharge_datetime = (datetime.strptime(admission_datetime, "%Y-%m-%d %H:%M:%S") +
                                  timedelta(days=hospitalization_days)).strftime('%Y-%m-%d %H:%M:%S')

            record_datetime = admission_datetime
            yy_collection_datetime = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            yy_etl_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

            # 医生信息
            attending_physician = fake.name()
            attending_physician_id = f"DOC{random.randint(100, 999)}"
            chief_physician = fake.name() if random.random() > 0.3 else None
            chief_physician_id = f"DOC{random.randint(100, 999)}" if chief_physician else None
            visit_doctor_name = attending_physician
            visit_doctor_no = attending_physician_id

            # 护士信息
            responsible_nurse = fake.name() if random.random() > 0.5 else None
            responsible_nurse_id = f"NUR{random.randint(100, 999)}" if responsible_nurse else None

            # 其他字段
            admission_type_code, admission_type_name = random.choice(ADMISSION_TYPES)
            discharge_reason = random.choice(DISCHARGE_REASONS)
            fee_type = random.choice(FEE_TYPES)
            hospitalization_times = str(random.randint(1, 5))
            special_flag = "肿瘤患者" if random.random() > 0.2 else None

            # 生成yy_record_md5 (模拟MD5)
            yy_record_md5 = f"md5_{random.getrandbits(128):032x}"
            yy_batch_time = day
            yy_record_batch_id = f"BATCH{day}_{i}"

            data = {
                'admission_bed_code': admission_bed_code,
                'admission_bed_name': admission_bed_name,
                'admission_datetime': admission_datetime,
                'admission_dept_code': dept_code,
                'admission_dept_name': dept_name,
                'admission_medical_team_code': medical_team_code,
                'admission_medical_team_name': medical_team_name,
                'admission_type_code': admission_type_code,
                'admission_type_name': admission_type_name,
                'admission_ward_code': ward_code,
                'admission_ward_name': ward_name,
                'attending_physician': attending_physician,
                'attending_physician_id': attending_physician_id,
                'chief_physician': chief_physician,
                'chief_physician_id': chief_physician_id,
                'discharge_bed_code': discharge_bed_code,
                'discharge_bed_name': discharge_bed_name,
                'discharge_datetime': discharge_datetime,
                'discharge_dept_code': discharge_dept_code,
                'discharge_dept_name': discharge_dept_name,
                'discharge_medical_team_code': medical_team_code,
                'discharge_medical_team_name': medical_team_name,
                'discharge_reason': discharge_reason,
                'discharge_ward_code': discharge_ward_code,
                'discharge_ward_name': discharge_ward_name,
                'extend_data1': None,
                'extend_data2': None,
                'fee_type': fee_type,
                'from_table': 'INPATIENT_RECORD',
                'from_yy_record_id': f"SOURCE_{random.randint(10000, 99999)}",
                'hospital_code': hospital_code,
                'hospital_name': hospital_name,
                'hospitalization_times': hospitalization_times,
                'inpatient_no': inpatient_no,
                'medical_record_no': medical_record_no,
                'name': name,
                'patient_id': patient_id,
                'patient_id_old': f"OLD_{patient_id}",
                'record_datetime': record_datetime,
                'record_status': 1,
                'record_update_datetime': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'responsible_nurse': responsible_nurse,
                'responsible_nurse_id': responsible_nurse_id,
                'special_flag': special_flag,
                'visit_card_no': f"VC{random.randint(100000, 999999)}",
                'visit_doctor_name': visit_doctor_name,
                'visit_doctor_no': visit_doctor_no,
                'visit_sn': visit_sn,
                'yy_collection_datetime': yy_collection_datetime,
                'yy_record_md5': yy_record_md5,
                'yy_upload_status': 0,
                'yy_etl_time': yy_etl_time,
                'yy_upload_time': None,
                'yy_batch_time': yy_batch_time,
                'yy_record_batch_id': yy_record_batch_id,
                'yy_backfill_time': None,
                'yy_backfill_status': None,
                'branch_code': branch_code,
                'branch_name': branch_name,
                'date_for_partition': admission_datetime
            }

            sql = _generate_sql('b03_1', data)
            sql_statements.append(sql)

    return sql_statements


def _generate_sql(table, data):
    columns = ', '.join([f'`{k}`' for k in data.keys()])
    values = []
    for v in data.values():
        if v is None:
            values.append('NULL')
        elif isinstance(v, (int, float)):
            values.append(str(v))
        else:
            escaped_value = str(v).replace("'", "''")
            values.append(f"'{escaped_value}'")
    return f"INSERT INTO `{table}` ({columns}) VALUES ({', '.join(values)});"


# ========== 执行生成 ==========
if __name__ == "__main__":
    records = generate_inpatient_records()
    # 写入数据库
    MysqlUtils_AI.insert_data_to_hub(records, 'b03_1')
    # 或写入 SQL 文件
    # FileUtil.generate_sql_file(records, "b03_1_inpatient_records.sql")